#!/usr/bin/env python
# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
# | 使用方法: python mddai.py <指令> [参数1,参数2,...]
# +----------------------------------------------------------------------
import sys
import time
import asyncio
from tortoise import Tortoise
from colorama import Fore, Style, init
from redis import asyncio as redis_asyncio
from config import get_settings
from common.utils.tools import ToolsUtil
from common.models.auth import AuthAdminModel

init()


class PrintOut:
    """ 打印工具 """
    @classmethod
    def success(cls, msg: str):
        print(Fore.GREEN + msg + Style.RESET_ALL)

    @classmethod
    def error(cls, msg: str):
        print(Fore.RED + msg + Style.RESET_ALL)


async def init_mysql():
    """ 初始化数据库连接 """
    # 数据库配置
    db_config = get_settings().DATABASES.get("connections") or {}
    mysql_config = db_config.get("mysql", {}).get("credentials")
    host = mysql_config["host"]
    port = mysql_config["port"]
    username = mysql_config["user"]
    password = mysql_config["password"]
    database = mysql_config["database"]

    # 初始化ORM
    await Tortoise.init(
        db_url=f"mysql://{username}:{password}@{host}:{port}/{database}",
        modules={"models": ["common.models.auth"]}
    )


async def init_redis():
    """ 初始化Redis连接 """
    # Redis配置
    redis_config = get_settings().REDIS
    host = redis_config.get("host", "127.0.0.1")
    port = redis_config.get("port", 6379)
    db = redis_config.get("db", 0)
    username = redis_config.get("username", "")
    password = redis_config.get("password", "")

    # 创建Redis
    redis_url = f"redis://{host}:{port}"
    redis = await redis_asyncio.from_url(
        redis_url,
        db=db,
        username=username,
        password=password,
        decode_responses=True
    )
    return redis


async def reset_password(params: list) -> bool:
    """ 重置管理员密码 """
    if len(params) < 0:
        PrintOut.error("错误: 请填写管理员账号和密码")
        PrintOut.error("示例: python mdd.py 1 account new_password")
        return False

    if len(params) < 2:
        PrintOut.error("错误: 请填写账号的新密码")
        PrintOut.error("示例: python mdd.py 1 account new_password")
        return False

    if len(params[1].strip()) < 6:
        PrintOut.error("错误: 新的密码长度不能少于 6 位数")
        return False

    if len(params[1].strip()) > 20:
        PrintOut.error("错误: 新的密码长度不能大于 20 位数")
        return False

    username = params[0].strip()
    password = params[1].strip()
    try:
        # 连接数据库
        await init_mysql()

        # 查管理员
        admin = await AuthAdminModel.filter(username=username, is_delete=0).first()
        if not admin:
            PrintOut.error(f"错误: 未找到名为 '{username}' 的管理员账户")
            return False

        try:
            # 新的密码
            new_salt = ToolsUtil.make_rand_char(6)
            new_password = ToolsUtil.make_md5_str(password, new_salt)

            # 修改密码
            admin.salt = new_salt
            admin.password = new_password
            admin.update_time = int(time.time())
            await admin.save()
            PrintOut.success(f"成功: 管理员 '{username}' 的密码已重置")
            return True
        except Exception as e:
            PrintOut.error(f"错误: 重置密码时发生异常 - {str(e)}")
            return False
    finally:
        await Tortoise.close_connections()


async def clear_redis_cache(params: list):
    if len(params) < 0:
        PrintOut.error("错误: 请指定清除场景 [system, login]")
        PrintOut.error("示例: python mdd.py 3 system")
        return False

    scene = str(params[0].strip())
    redis = None
    try:
        redis_config = get_settings().REDIS
        redis_prefix = str(redis_config.get("prefix", "wait:"))
        redis = await init_redis()

        if scene == "system":
            key = str(redis_prefix + "sys:config")
            await redis.delete(key)
            PrintOut.success(f"提示: 准备开始删除缓存键名 {key} ~")
            PrintOut.success(f"成功: 已清理项目的系统缓存 ~")
        elif scene == "login":
            cursor = 0
            number = 0
            match = str(redis_prefix + "login:*")
            while True:
                cursor, keys = await redis.scan(cursor=cursor, match=match)
                for key in keys:
                    await redis.delete(key)
                    number += 1
                if cursor == 0:
                    break
            PrintOut.success(f"提示: 开始扫描缓存键 {match} ~")
            PrintOut.success(f"成功: 已清理登录数据 {number} 条 ~")
    except Exception as e:
        PrintOut.error(f"错误: 清除缓存时发生异常 - {str(e)}")
    finally:
        if redis:
            await redis.aclose()


async def clear_login_limit(params: list):
    if len(params) < 0:
        PrintOut.error("错误: 请输入要解除限制的账号")
        PrintOut.error("示例: python mdd.py 4 account")
        return False

    account = params[0].strip()
    redis = None
    try:
        redis_config = get_settings().REDIS
        redis_prefix = str(redis_config.get("prefix", "wait:"))
        redis = await init_redis()

        key = str(redis_prefix + "login:lock_" + account)
        if not await redis.exists(key):
            PrintOut.error(f"提示: 未找到账号为 '{account}' 的登录限制 ~")
            return False

        ttl = await redis.ttl(key)
        await redis.delete(key)
        PrintOut.success(f"提示: '{account}' 账号限制时间剩余 {ttl} 秒 ~")
        PrintOut.success(f"成功: '{account}' 账号限制已被清除 ~")
    except Exception as e:
        PrintOut.error(f"错误: 清除限制时发生异常 - {str(e)}")
    finally:
        if redis:
            await redis.aclose()


async def main():
    # 检查命令行参数
    params = sys.argv[1:]
    if len(params) == 0:
        old, new = "           ", ""
        return print("""
            ===============指令工具箱==============
            (1) 修改后台密码   (2) 获取当前版本
            (3) 清洗项目缓存   (4) 登录限制清除
            ======================================
            请携带命令编号如: python command.py 1
        """.replace(old, new))

    instruct = params[0]
    if instruct == "1":
        await reset_password(params[1:])
    elif instruct == "2":
        version: str = get_settings().VERSION
        PrintOut.success(f"当前系统版本为: v{version}")
    elif instruct == "3":
        await clear_redis_cache(params[1:])
    elif instruct == "4":
        await clear_login_limit(params[1:])
    else:
        PrintOut.error(f"错误: 没有找到对应的指令编号 '{instruct}'")


if __name__ == "__main__":
    asyncio.run(main())
